{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## VAE\n",
    "VAE（Variational Autoencoder，变分自编码器）是一种生成式深度学习模型，属于无监督学习方法，可用于图片生成。\n",
    "\n",
    "VAE论文：https://arxiv.org/abs/1312.6114"
   ],
   "id": "1b1d4c3fbbc6827b"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "## Latent Image--->潜层空间向量--->加噪声--->噪声--->",
   "id": "90f4b383f631cbd3"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-26T08:52:06.799564Z",
     "start_time": "2025-08-26T08:52:06.232438Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class VAE(nn.Module):\n",
    "    def __init__(self, in_channels=1, latent_channels=4):\n",
    "        super().__init__()\n",
    "        self.latent_channels = latent_channels\n",
    "        # Encoder\n",
    "        self.encoder = nn.Sequential(\n",
    "            nn.Conv2d(in_channels, 32, kernel_size=4, stride=2, padding=1),  # 32x32 -> 16x16\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1),  # 16x16 -> 8x8\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),  # 8x8 -> 4x4\n",
    "            nn.ReLU(),\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(128 * 4 * 4, 256)\n",
    "        )\n",
    "        self.fc_mu = nn.Linear(256, latent_channels * 4 * 4)\n",
    "        self.fc_log_var = nn.Linear(256, latent_channels * 4 * 4)\n",
    "\n",
    "        # Decoder\n",
    "        self.decoder_input = nn.Linear(latent_channels * 4 * 4, 128 * 4 * 4)\n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.Unflatten(1, (128, 4, 4)),\n",
    "            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),  # 4x4 -> 8x8\n",
    "            nn.ReLU(),\n",
    "            nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),  # 8x8 -> 16x16\n",
    "            nn.ReLU(),\n",
    "            nn.ConvTranspose2d(32, in_channels, kernel_size=4, stride=2, padding=1),  # 16x16 -> 32x32\n",
    "            nn.Tanh()  # Output values between -1 and 1\n",
    "        )\n",
    "\n",
    "    def encode(self, x):\n",
    "        h = self.encoder(x)\n",
    "        mu = self.fc_mu(h)\n",
    "        log_var = self.fc_log_var(h)\n",
    "        return mu, log_var\n",
    "\n",
    "    def reparameterize(self, mu, log_var):\n",
    "        std = torch.exp(0.5 * log_var)\n",
    "        eps = torch.randn_like(std)\n",
    "        return mu + eps * std\n",
    "\n",
    "    def decode(self, z):\n",
    "        h = self.decoder_input(z)\n",
    "        return self.decoder(h)\n",
    "\n",
    "    def forward(self, x):\n",
    "        mu, log_var = self.encode(x)\n",
    "        z = self.reparameterize(mu, log_var)\n",
    "        return self.decode(z), mu, log_var\n",
    "\n",
    "    def get_latent(self, x):\n",
    "        mu, log_var = self.encode(x)\n",
    "        z = self.reparameterize(mu, log_var)\n",
    "        return z.view(-1, self.latent_channels, 4, 4)"
   ],
   "id": "b98c5adcc3bd70f0",
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-26T08:52:07.411798Z",
     "start_time": "2025-08-26T08:52:06.886804Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize(32),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.5], [0.5])  # Normalize to [-1, 1]\n",
    "])\n",
    "dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)\n",
    "vae_dataloader = DataLoader(dataset, batch_size=256, shuffle=True, num_workers=4)"
   ],
   "id": "c603a8f881ce9a2e",
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-26T08:53:13.120317Z",
     "start_time": "2025-08-26T08:52:07.420965Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from torchvision.utils import save_image\n",
    "import os\n",
    "import torch.nn.functional as F\n",
    "\n",
    "IN_CHANNELS = 1\n",
    "LATENT_CHANNELS = 4\n",
    "# VAE_EPOCHS = 25\n",
    "VAE_EPOCHS = 1\n",
    "\n",
    "os.makedirs(\"vae_samples\", exist_ok=True)\n",
    "os.makedirs(\"final_weights\", exist_ok=True)\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "vae = VAE(in_channels=IN_CHANNELS, latent_channels=LATENT_CHANNELS).to(device)\n",
    "vae_optimizer = torch.optim.Adam(vae.parameters(), lr=1e-3)\n",
    "\n",
    "def vae_loss_function(recon_x, x, mu, log_var):\n",
    "    recon_loss = F.mse_loss(recon_x, x, reduction='sum')\n",
    "    kld = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())\n",
    "    return recon_loss + kld\n",
    "\n",
    "\n",
    "for epoch in range(VAE_EPOCHS):\n",
    "    vae.train()\n",
    "    for images, _ in vae_dataloader:\n",
    "        images = images.to(device)\n",
    "        vae_optimizer.zero_grad()\n",
    "        recon_images, mu, log_var = vae(images)\n",
    "        loss = vae_loss_function(recon_images, images, mu, log_var)\n",
    "        loss.backward()\n",
    "        vae_optimizer.step()\n",
    "\n",
    "    if (epoch + 1) % 5 == 0:\n",
    "        vae.eval()\n",
    "        with torch.no_grad():\n",
    "            fixed_images, _ = next(iter(vae_dataloader))\n",
    "            recon_fixed, _, _ = vae(fixed_images.to(device))\n",
    "            comparison = torch.cat([fixed_images[:8], recon_fixed.cpu()[:8]])\n",
    "            save_image(comparison, f\"vae_samples/epoch_{epoch + 1}.png\", nrow=8, normalize=True)\n",
    "\n",
    "vae_path = \"final_weights/mnist_vae.pth\"\n",
    "torch.save(vae.state_dict(), vae_path)\n",
    "print(f\"VAE模型训练完成，模型保存在了{vae_path}\")"
   ],
   "id": "d1632b51bd955f87",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VAE模型训练完成，模型保存在了final_weights/mnist_vae.pth\n"
     ]
    }
   ],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-26T08:53:13.142373Z",
     "start_time": "2025-08-26T08:53:13.131025Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 从本地加载VAE模型\n",
    "vae_path = \"final_weights/mnist_vae.pth\"\n",
    "vae = VAE(in_channels=IN_CHANNELS, latent_channels=LATENT_CHANNELS).to(device)\n",
    "vae.load_state_dict(torch.load(vae_path))\n",
    "vae.eval()\n",
    "vae.requires_grad_(False)\n",
    "\n",
    "\n",
    "# 随机生成图片\n",
    "fixed_images = torch.randn(8, IN_CHANNELS, 32, 32).to(device)\n",
    "result = vae.get_latent(fixed_images.to(device))\n",
    "print(result.shape)"
   ],
   "id": "b9f03180ffa9614a",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([8, 4, 4, 4])\n"
     ]
    }
   ],
   "execution_count": 4
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
